import os
import json
import torch

import numpy as np
import torch.nn.functional as F
import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET

from scipy.ndimage import gaussian_filter


def get_map_from_points(points, shape, sigma=8):
    """
    Generate the density map based on the dot annotations provided by the points.
    """
    label = torch.tensor(points).float()

    height, width = shape
    density_map = torch.zeros((1, height, width), dtype=torch.float32)

    if len(label) > 0:
        assert len(label.shape) == 2 and label.shape[1] == 2, f"label should be a Nx2 tensor, got {label.shape}."
        label_ = label.long()
        label_[:, 0] = label_[:, 0].clamp(min=0, max=width - 1)
        label_[:, 1] = label_[:, 1].clamp(min=0, max=height - 1)
        density_map[0, label_[:, 1], label_[:, 0]] = 1.0

    if sigma is not None:
        assert sigma > 0, f"sigma should be positive if not None, got {sigma}."
        density_map = torch.from_numpy(gaussian_filter(density_map, sigma=sigma))

    return density_map.squeeze()


def get_json_points(file_path):
    with open(file_path, 'r') as file:
        points = json.load(file)
    points = points['points']
    return points


def get_xml_points(label):
    tree = ET.parse(label)
    root = tree.getroot()
    points = [[int(obj.find('point').find('x').text), int(obj.find('point').find('y').text)]for obj in root.findall('object')]
    return points


def get_map_num(pt_file, shape):
    """
    shape = (image_height, image_width)
    """
    dens = torch.load(pt_file).cpu()
    count = dens.sum().item()
    dens_sum = torch.sum(dens, dim=(-1, -2))
    
    up_dens = F.interpolate(dens, size=shape, mode="bilinear")
    up_dens_sum = torch.sum(up_dens, dim=(-1, -2))
    up_scale = torch.nan_to_num(up_dens_sum / dens_sum, nan=0.0, posinf=0.0, neginf=0.0)
    
    up_dens *= up_scale
    up_dens = up_dens.squeeze()
    
    return up_dens, round(count, 2)


def get_fidtm_map_num(pt_file):
    pt = torch.load(pt_file).cpu()

    # map
    dens = pt.numpy()
    dens[dens < 0] = 0
    dens = 255 * dens / np.max(dens)
    dens = dens.squeeze()
    dens = torch.from_numpy(dens)

    # count
    pt_max = torch.max(pt).item()

    keep = torch.nn.functional.max_pool2d(pt, (3, 3), stride=1, padding=1)
    keep = (keep == pt).float()
    pt = keep * pt

    '''set the pixel valur of local maxima as 1 for counting'''
    pt[pt < 100.0 / 255.0 * pt_max] = 0
    pt[pt > 0] = 1

    ''' negative sample'''
    if pt_max < 0.1:
        pt = pt * 0

    count = pt.sum().item()

    return dens, round(count, 2)


def save_density(density, save_file):
    plt.imshow(density, cmap="jet")
    plt.axis('off')
    plt.savefig(save_file, bbox_inches='tight', pad_inches=0)
    plt.close()


def batch_save_density(result_root, names, sahpe, save_root, method):
    """
    shape = (image_height, image_width)
    """
    for name in names:
        pt_file = os.path.join(result_root, name+'.pt')
        density, count = get_map_num(pt_file, sahpe)
        
        save_name = f"{name}_{method}_{round(count, 2)}.png"
        save_file = os.path.join(save_root, save_name)
        save_density(density, save_file)
        print(f"[INFO]batch_save_density: {save_file}")
        
        
        
if __name__ == '__main__':
    save_root = './DroneRGBT_density'
    if not os.path.exists(save_root):
        os.makedirs(save_root)
    
    names = ['1047_RGB', '131_RGB', '1599_RGB', '605_RGB', '860_RGB']
    result_root = '/data/store1/nzd/rgbtcc_work/rgbtcc_works/other_methods/MIANet/map_DroneRGBT/'
    batch_save_density(result_root, names, (512, 640), save_root, 'MIANet')
    
    save_root = './RGBTCC_density'
    if not os.path.exists(save_root):
        os.makedirs(save_root)
    
    names = ['1325_RGB', '1684_RGB', '1729_RGB', '2227_RGB', '3401_RGB']
    result_root = '/data/store1/nzd/rgbtcc_work/rgbtcc_works/other_methods/MIANet/map_RGBTCC/'
    batch_save_density(result_root, names, (480, 640), save_root, 'MIANet')